'''
Created on Nov 15, 2009

@author: mkiyer
'''
'''
Created on Sep 18, 2009

@author: mkiyer
'''

from optparse import OptionParser
from veggie.coverage.visualization.ucsc import get_bedgraph_track_header, \
    gen_bedgraph_track_data
from veggie.gene.gene import parse_refseq_genes
from veggie.sample.sampledb2 import get_sampledb
from veggie.sample.samplegroup import parse_samplegroups_xml
import collections
import h5py
import logging
import matplotlib.pyplot as plt
import numpy as np
import sys
import os
import veggie.coverage.covdb as covdb
import veggie.coverage.visualization.plots as oncoplot


def read_genes_file(fhd, column=0):
    gene_lookup = {}
    for g in parse_refseq_genes():
        gene_lookup[g.symbol] = g
        gene_lookup[g.acc] = g
    for line in fhd:
        line = line.strip()
        if line is None:
            continue
        if line.startswith('#'):
            continue        
        fields = line.split('\t')
        g = gene_lookup.get(fields[column], None)
        if g is None:
            logging.debug('Could not find gene %s, skipping...\n' % fields[column])
            yield None, fields
        else:
            yield g, fields
        
if __name__ == '__main__':
    results_path='./results'
    # configure logging
    logging.basicConfig(level=logging.DEBUG)
    # parse command line
    optionparser = OptionParser("usage: %prog [options] <regions.bed>")
    optionparser.add_option("--covdb", dest="covdb",
                            help="path to coverage database")
    optionparser.add_option("-s", "--samples", dest="samples",
                            help="sample group XML file")
    optionparser.add_option("-o", "--output", dest="outfile",
                            default=None,
                            help="output file [default: %default]")
    optionparser.add_option("--col", dest="gene_col", type="int",
                            default=0)
    optionparser.add_option("--lanes", action="store_false", dest="pool_lanes",
                            default=True,
                            help="show individual lane results rather "
                            "than pooling")

    (options, args) = optionparser.parse_args()
    if len(args) == 0:
        optionparser.error("no BED file specified")

    # open coverage db
    h5file = h5py.File(options.covdb, 'r')
    # generate SampleGroupCoverage objects that pair samples with
    # coverage data
    sample_groups = [covdb.SampleGroupCoverage(sgroup, h5file) for sgroup in 
                     list(parse_samplegroups_xml(options.samples))]
    # for each sample generate an exon expression matrix
    for sgroup in sample_groups:
        for sample_name, sample_dsets in sgroup.iteritems():
            # open a file for the exon expressions
            outfhd = open(os.path.join(results_path, "%s_exon_expr.txt" % (sample_name)), 'w')
            # sample properties header
            outfhd.write('Sample Group\t%s\n' % sgroup.name)
            outfhd.write('Sample Name\t%s\n' % sample_name)
            # write all properties in param_schema
            for prop in sorted(get_sampledb().param_schema.keys()):
                outfhd.write('%s' % prop)
                if prop in get_sampledb().params[sample_name]:
                    outfhd.write('\t%s' % get_sampledb().params[sample_name][prop])
                else:
                    outfhd.write('\tNone')
                outfhd.write('\n')
            outfhd.write('gene\texon_0\texon_1\texon_2\t...\n')
            
            # read input genes file
            for res in read_genes_file(open(args[0]), options.gene_col):
                g, fields = res
                # print fields first
                outfhd.write('%s' % (fields[0]))
                for f in fields[1:]:
                    outfhd.write('\t%s' % f)
                if g is None:
                    outfhd.write('\tNone\n')
                    continue
                # now print exon coverage
                for e in g.exons:
                    chrom = g.chrom
                    start, end = e[0], e[1]
                    if options.pool_lanes:
                        # here we have the interval, the sample name, and the
                        # datasets that comprise the sample.  we can now 
                        # compute coverage
                        avgcov = covdb.calc_avg_cov_rpkm(chrom, start, end, sample_dsets)
                        outfhd.write('\t%.3f' % avgcov)
                    else:
                        logging.debug('separate lanes disabled')
                        sys.exit(0)
                        rpkms = [('%.2f' % covdb.calc_rpkm(chrom, start, end, dset)) 
                                 for dset in sample_dsets]
                        outfhd.write('\t')
                        outfhd.write(','.join(rpkms))
                outfhd.write('\n')
            # close exon expressions file
            outfhd.close()    
    h5file.close()
